# -*- coding: utf-8 -*-
"""
数据处理脚本

处理数据集，将其转换为可用于训练神经网络的格式，并划分训练集和测试集。
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import LabelEncoder, StandardScaler, OneHotEncoder
from sklearn.model_selection import train_test_split
from sklearn.impute import SimpleImputer
import pickle
import os

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 创建输出目录
if not os.path.exists('output'):
    os.makedirs('output')

# 创建模型目录
if not os.path.exists('models'):
    os.makedirs('models')

# 创建图表目录
if not os.path.exists('figures'):
    os.makedirs('figures')


def load_data(file_path='数据集0727.xlsx'):
    """
    加载数据集
    """
    print(f"正在加载数据集: {file_path}...")
    data = pd.read_excel(file_path)
    print(f"数据集加载完成，形状: {data.shape}")
    return data


def clean_data(data):
    """
    清洗数据，处理缺失值、异常值等
    
    按照以下要求处理数据：
    1. 输出列的所有值必须存在，如果值有不存在的，则删除该行
    2. 输入列中必须有值的列有：高分子基体（缩写）、无机添加剂（缩写），其他的（其余材料、加工方法）可以为空
       如果必须有值的列中有'/'或空值，直接删除该行
    3. 对于可以为空的输入列，保留缺失值
    """
    print("开始数据清洗...")
    
    # 只保留需要的列
    input_cols = ['高分子基体（缩写）', '无机添加剂（缩写）', '其余材料', '加工方法']
    required_input_cols = ['高分子基体（缩写）', '无机添加剂（缩写）']  # 必须有值的输入列
    optional_input_cols = ['其余材料', '加工方法']  # 可以为空的输入列
    output_cols = ['太阳反射率', '红外发射率']
    
    # 检查是否所有列都存在
    for col in input_cols + output_cols:
        if col not in data.columns:
            raise ValueError(f"数据集中缺少列: {col}")
    
    # 提取需要的列
    df = data[input_cols + output_cols].copy()
    print(f"原始数据集形状: {df.shape}")
    
    # 处理特殊字符、空格和将'/'标记为空值（在删除行之前先处理）
    for col in input_cols:
        if df[col].dtype == 'object':
            # 替换特殊字符和多余空格
            df[col] = df[col].astype(str).str.replace('\u200b', '')
            df[col] = df[col].str.strip()
            
            # 将'nan'字符串、空字符串和'/'标记为空值
            df[col] = df[col].replace(['nan', '/', ''], np.nan)
    
    # 1. 删除输出列中有缺失值的行
    missing_output = df[output_cols].isnull().any(axis=1).sum()
    if missing_output > 0:
        print(f"删除 {missing_output} 行输出特征缺失值")
    df = df.dropna(subset=output_cols)
    
    # 2. 删除必须有值的输入列中有缺失值的行
    missing_required = df[required_input_cols].isnull().any(axis=1).sum()
    if missing_required > 0:
        print(f"删除 {missing_required} 行必须有值的输入特征缺失值")
    df = df.dropna(subset=required_input_cols)
    print(f"清洗后的数据集形状: {df.shape}")
    
    # 3. 对于可选输入列，保留缺失值（不需要额外处理）
    for col in optional_input_cols:
        missing_count = df[col].isnull().sum()
        if missing_count > 0:
            print(f"列 '{col}' 有 {missing_count} 个缺失值，将在预处理阶段处理")
    
    # 处理输出特征，将百分数和小数统一为小数格式，并保留四位小数
    for col in output_cols:
        # 检查列中的非空值
        non_null_mask = ~df[col].isnull()
        if non_null_mask.any():
            # 将列转换为字符串以便处理
            df.loc[non_null_mask, col] = df.loc[non_null_mask, col].astype(str)
            
            # 清理特殊字符（制表符、零宽空格等）
            df.loc[non_null_mask, col] = df.loc[non_null_mask, col].str.replace('\t', '')
            df.loc[non_null_mask, col] = df.loc[non_null_mask, col].str.replace('\u200b', '')
            df.loc[non_null_mask, col] = df.loc[non_null_mask, col].str.strip()
            
            # 处理百分数格式（如"96.7%"）
            percent_mask = df[col].str.contains('%', na=False)
            if percent_mask.any():
                # 移除百分号并转换为小数
                df.loc[percent_mask, col] = df.loc[percent_mask, col].str.replace('%', '')
                # 将百分数转换为小数（除以100）
                df.loc[percent_mask, col] = pd.to_numeric(df.loc[percent_mask, col], errors='coerce') / 100
            
            # 将剩余的转换为数值类型
            df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # 确保列是数值类型后，保留四位小数
        if pd.api.types.is_numeric_dtype(df[col]):
            # 使用round函数保留4位小数
            df[col] = df[col].round(4)
            # 将数值转换为字符串，确保显示4位小数
            df[col] = df[col].apply(lambda x: '{:.4f}'.format(x) if pd.notnull(x) else x)
            # 再转回浮点数
            df[col] = pd.to_numeric(df[col])
    
    # 再次删除可能因转换而产生的缺失值
    numeric_missing = df[output_cols].isnull().any(axis=1).sum()
    if numeric_missing > 0:
        print(f"删除 {numeric_missing} 行因数值转换产生的缺失值")
    df = df.dropna(subset=output_cols)
    
    # 检查输出特征的范围是否在0-1之间
    for col in output_cols:
        out_of_range = ((df[col] < 0) | (df[col] > 1)).sum()
        if out_of_range > 0:
            print(f"警告: 列 '{col}' 中有 {out_of_range} 个值不在0-1范围内，将进行裁剪")
            df[col] = df[col].clip(0, 1)
    
    print("数据清洗完成")
    return df


def preprocess_data(df):
    """
    数据预处理，包括编码分类变量、标准化等
    
    处理逻辑：
    1. 对于必须有值的输入列（高分子基体、无机添加剂），直接进行OneHotEncoder编码
    2. 对于可以为空的输入列（其余材料、加工方法），将缺失值标记为'缺失'类别，然后进行编码
       确保缺失值作为一个独立的特征值被考虑在内
    3. 所有分类变量都使用OneHotEncoder进行编码
    """
    print("开始数据预处理...")
    
    input_cols = ['高分子基体（缩写）', '无机添加剂（缩写）', '其余材料', '加工方法']
    required_input_cols = ['高分子基体（缩写）', '无机添加剂（缩写）']  # 必须有值的输入列
    optional_input_cols = ['其余材料', '加工方法']  # 可以为空的输入列
    output_cols = ['太阳反射率', '红外发射率']
    
    # 创建特征编码器字典
    encoders = {}
    
    # 处理分类特征 - 使用OneHotEncoder
    X_encoded_parts = []
    feature_names = []
    
    for col in input_cols:
        print(f"处理特征: {col}")
        
        # 创建一个临时DataFrame用于编码
        temp_df = df[[col]].copy()
        
        # 处理缺失值
        if col in optional_input_cols:
            # 对于可选输入列，将缺失值标记为'缺失'类别
            missing_count = temp_df[col].isnull().sum()
            if missing_count > 0:
                print(f"  列 '{col}' 有 {missing_count} 个缺失值，将标记为'缺失'类别并作为独立特征值处理")
                temp_df[col] = temp_df[col].fillna('缺失')
                # 确保'缺失'被视为一个有效的类别，并在OneHotEncoder中被正确编码
                # 这样缺失值就会作为一个独立的特征值被考虑在内
        else:
            # 对于必须有值的列，确认没有缺失值
            missing_count = temp_df[col].isnull().sum()
            if missing_count > 0:
                print(f"  警告: 必须有值的列 '{col}' 仍有 {missing_count} 个缺失值")
                # 这种情况不应该发生，因为我们已经在clean_data中删除了这些行
                # 但为了健壮性，我们再次检查并删除
                raise ValueError(f"列 '{col}' 中存在缺失值，请检查数据清洗步骤")
        
        # 创建OneHotEncoder，设置handle_unknown='ignore'以处理未知类别
        encoder = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
        
        # 拟合并转换
        encoded = encoder.fit_transform(temp_df)
        
        # 获取特征名称
        categories = encoder.categories_[0]
        col_names = [f"{col}_{cat}" for cat in categories]
        feature_names.extend(col_names)
        
        # 保存编码器
        encoders[col] = encoder
        
        # 添加到编码部分列表
        X_encoded_parts.append(encoded)
    
    # 合并所有编码后的特征
    X_encoded = np.hstack(X_encoded_parts)
    
    # 创建DataFrame以便查看
    X_encoded_df = pd.DataFrame(X_encoded, columns=feature_names)
    print(f"编码后的特征形状: {X_encoded.shape}")
    
    # 提取输出特征
    y = df[output_cols].values
    
    # 保存特征名称和编码器
    with open('output/feature_names.pkl', 'wb') as f:
        pickle.dump(feature_names, f)
    
    with open('output/encoders.pkl', 'wb') as f:
        pickle.dump(encoders, f)
    
    print("数据预处理完成，特征名称和编码器已保存")
    return X_encoded, y, feature_names, encoders


def split_dataset(X, y, test_size=0.2, random_state=42):
    """
    划分训练集和测试集
    
    参数:
        X: 特征矩阵
        y: 目标变量
        test_size: 测试集比例，默认0.2
        random_state: 随机种子，默认42
    
    返回:
        X_train, X_test, y_train, y_test: 训练集和测试集的特征和目标变量
    """
    print(f"划分数据集，测试集比例: {test_size}")
    
    # 确保输入数据的有效性
    if X.shape[0] != y.shape[0]:
        raise ValueError(f"特征数量 ({X.shape[0]}) 与目标变量数量 ({y.shape[0]}) 不匹配")
    
    if X.shape[0] < 10:
        print("警告: 样本数量过少，可能导致划分不均衡")
    
    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=None
    )
    
    print(f"训练集形状: X_train {X_train.shape}, y_train {y_train.shape}")
    print(f"测试集形状: X_test {X_test.shape}, y_test {y_test.shape}")
    
    return X_train, X_test, y_train, y_test


def visualize_data(df, X_encoded, feature_names):
    """
    可视化数据分布和特征关系
    
    参数:
        df: 原始数据DataFrame
        X_encoded: 编码后的特征矩阵
        feature_names: 特征名称列表
    """
    print("开始数据可视化...")
    
    # 确保输出目录存在
    if not os.path.exists('figures'):
        os.makedirs('figures')
        print("创建figures目录用于保存可视化结果")
    
    # 检查输入数据的有效性
    if len(df) == 0 or X_encoded.shape[0] == 0:
        print("警告: 输入数据为空，无法进行可视化")
        return
    
    if '太阳反射率' not in df.columns or '红外发射率' not in df.columns:
        print("警告: 输入数据缺少必要的输出列，无法进行完整可视化")
        return
    
    # 创建一个图形，包含多个子图
    plt.figure(figsize=(15, 12))
    
    # 1. 太阳反射率分布
    plt.subplot(2, 2, 1)
    sns.histplot(df['太阳反射率'], kde=True)
    plt.title('太阳反射率分布')
    plt.xlabel('太阳反射率')
    plt.ylabel('频数')
    
    # 2. 红外发射率分布
    plt.subplot(2, 2, 2)
    sns.histplot(df['红外发射率'], kde=True)
    plt.title('红外发射率分布')
    plt.xlabel('红外发射率')
    plt.ylabel('频数')
    
    # 3. 太阳反射率与红外发射率的关系
    plt.subplot(2, 2, 3)
    sns.scatterplot(x='太阳反射率', y='红外发射率', data=df)
    plt.title('太阳反射率与红外发射率的关系')
    plt.xlabel('太阳反射率')
    plt.ylabel('红外发射率')
    
    # 4. 特征重要性（使用相关性作为简单指标）
    plt.subplot(2, 2, 4)
    
    try:
        # 将编码后的特征与输出合并
        X_encoded_df = pd.DataFrame(X_encoded, columns=feature_names)
        X_encoded_df['太阳反射率'] = df['太阳反射率'].values
        X_encoded_df['红外发射率'] = df['红外发射率'].values
        
        # 计算与输出的相关性
        corr_solar = X_encoded_df.corr()['太阳反射率'].drop(['太阳反射率', '红外发射率'])
        corr_ir = X_encoded_df.corr()['红外发射率'].drop(['太阳反射率', '红外发射率'])
    except Exception as e:
        print(f"计算特征相关性时出错: {e}")
        return
    
    try:
        # 合并相关性并取绝对值
        corr_combined = pd.DataFrame({
            '太阳反射率相关性': corr_solar.abs(),
            '红外发射率相关性': corr_ir.abs()
        })
        
        # 检查是否有足够的特征进行可视化
        if len(corr_combined) == 0:
            print("警告: 没有足够的特征进行相关性分析")
            return
        
        # 取前N个特征（如果特征数量不足10个，则取全部）
        n_features = min(10, len(corr_combined))
        top_features = corr_combined.mean(axis=1).sort_values(ascending=False).head(n_features).index
        corr_plot = corr_combined.loc[top_features]
        
        # 绘制条形图
        ax = corr_plot.plot(kind='barh')
        plt.title('特征与输出的相关性（前{}个特征）'.format(n_features))
        plt.xlabel('相关性系数（绝对值）')
        plt.ylabel('特征')
        plt.tight_layout()
        
        # 保存图形
        plt.savefig('figures/data_visualization.png', dpi=300, bbox_inches='tight')
        print("数据可视化完成，图表已保存到 figures/data_visualization.png")
    except Exception as e:
        print(f"绘制相关性图表时出错: {e}")
        return


def save_processed_data(X_train, X_test, y_train, y_test):
    """
    保存处理后的数据
    
    参数:
        X_train: 训练集特征
        X_test: 测试集特征
        y_train: 训练集目标变量
        y_test: 测试集目标变量
    """
    print("保存处理后的数据...")
    
    # 确保输出目录存在
    if not os.path.exists('output'):
        os.makedirs('output')
        print("创建output目录用于保存处理后的数据")
    
    # 检查输入数据的有效性
    if X_train.shape[0] == 0 or X_test.shape[0] == 0:
        print("警告: 训练集或测试集为空，无法保存")
        return
    
    if X_train.shape[0] != y_train.shape[0] or X_test.shape[0] != y_test.shape[0]:
        print("警告: 特征数量与目标变量数量不匹配，无法保存")
        return
    
    try:
        # 保存数据
        np.save('output/X_train.npy', X_train)
        np.save('output/X_test.npy', X_test)
        np.save('output/y_train.npy', y_train)
        np.save('output/y_test.npy', y_test)
        print("数据已保存到 output/ 目录")
    except Exception as e:
        print(f"保存数据时出错: {e}")
        return


def main(file_path='数据集0727.xlsx', test_size=0.1, random_state=66):
    """
    主函数，执行整个数据处理流程
    
    参数:
        file_path: 数据文件路径，默认为'数据集0727.xlsx'
        test_size: 测试集比例，默认0.2
        random_state: 随机种子，默认42
    """
    try:
        print("\n===== 开始数据处理流程 =====")
        
        # 1. 加载数据
        data = load_data(file_path)
        if data.empty:
            print("错误: 加载的数据为空，无法继续处理")
            return
        
        # 2. 清洗数据 - 处理缺失值和异常值
        df = clean_data(data)
        if df.empty:
            print("错误: 清洗后的数据为空，无法继续处理")
            return
        
        # 3. 预处理数据 - 编码分类变量（此步骤已保存特征名称和编码器）
        X_encoded, y, feature_names, encoders = preprocess_data(df)
        if X_encoded.shape[0] == 0 or len(feature_names) == 0:
            print("错误: 预处理后的数据为空，无法继续处理")
            return
        
        # 4. 划分数据集 - 训练集和测试集
        X_train, X_test, y_train, y_test = split_dataset(X_encoded, y, test_size, random_state)
        
        # 5. 可视化数据 - 分析特征与输出的相关性
        visualize_data(df, X_encoded, feature_names)
        
        # 6. 保存处理后的数据 - 用于模型训练
        save_processed_data(X_train, X_test, y_train, y_test)
        
        # 7. 保存清洗后的数据 - 用于后续分析
        try:
            # 确保太阳反射率和红外发射率列保留4位小数
            output_cols = ['太阳反射率', '红外发射率']
            
            # 创建一个新的DataFrame用于保存，以避免修改原始DataFrame
            df_to_save = df.copy()
            
            for col in output_cols:
                if col in df_to_save.columns and pd.api.types.is_numeric_dtype(df_to_save[col]):
                    # 将数值转换为字符串，确保显示4位小数
                    df_to_save[col] = df_to_save[col].apply(lambda x: '{:.4f}'.format(x) if pd.notnull(x) else '')
            
            # 保存到CSV，不使用float_format，因为已经将列转换为字符串
            df_to_save.to_csv('output/cleaned_data.csv', index=False, encoding='utf-8-sig')
            print("清洗后的数据已保存到 output/cleaned_data.csv")
            
            # 为了确保在后续处理中使用的是原始数值类型，我们不修改原始df
        except Exception as e:
            print(f"保存清洗后的数据时出错: {e}")
        
        # 8. 输出处理结果摘要
        print("\n===== 数据处理完成 =====")
        print(f"特征数量: {len(feature_names)}")
        print(f"样本总数: {len(df)}")
        print(f"训练样本: {len(X_train)} ({100 * len(X_train) / len(df):.1f}%)")
        print(f"测试样本: {len(X_test)} ({100 * len(X_test) / len(df):.1f}%)")
        print("\n处理后的数据、特征名称和编码器已保存到 output/ 目录")
        print("可视化结果已保存到 figures/ 目录")
        
    except Exception as e:
        print(f"\n错误: 数据处理过程中发生异常: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()